"""
https://paperswithcode.com/method/layer-normalization
https://github.com/CyberZHG/torch-layer-normalization/blob/89f405b60f53f85da6f03fe685c190ef394ce50c/torch_layer_normalization/layer_normalization.py#L8
"""

import torch

B, S, E = 4, 512, 64
ln = torch.nn.LayerNorm((E, ))
print('LayerNorm', ln)

input = torch.randn(B, S, E)
for b in range(B):
    input[b] *= 5.0
    input[b] += 3.0
print('input', input.shape)
print('std', input[0, 0, :].std())
print('mean', input[0, 0, :].mean())
output = ln(input)
print('output', output.shape)
print('std', output[0, 0, :].std())
print('mean', output[0, 0, :].mean())
